#define JUGGLE_VERSION "0.1.2"

#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <string.h>
#include <math.h>
#include <argp.h>

#include "config.h"
#include "utils.h"

const char* argp_program_version = JUGGLE_VERSION;

void add_cid(int cid, config_t* config) {
	if (cid < 0 || cid >= MAX_CPUS) {
		FAIL("CPU id %d in map out of range: must be >= 0 and < %d", cid, MAX_CPUS);
	}
	for (int i = 0; i < config->num_cpus; i++) {
		if (config->cids[i] == cid) return;
	}
	config->cids[config->num_cpus] = cid;
	config->num_cpus++;
	if (config->num_cpus > MAX_CPUS) {
		FAIL("Only support %d cores at the moment, not %d", MAX_CPUS, config->num_cpus);
	}
	//printf("%d\n", cid);
}

static error_t parse_opt(int key, char* optarg, struct argp_state *state) {
	config_t* config = state->input;
	char* flag_token;
	switch (key) {
	case 'c':
		config->num_cpus = 0;
		char* token = strtok(optarg, ",");
		while (token) {
			char* c1 = strchr(token, '-');
			if (c1) {
				int cid1 = atoi(strsep(&token, "-"));
				int cid2 = atoi(c1 + 1);
				for (int i = cid1; i <= cid2; i++) add_cid(i, config);
			} else {
				add_cid(atoi(token), config);
			}
			token = strtok(NULL, ",");
		}
		break;
	case 'n': config->num_tasks = atoi(optarg); break;
	case 'i': config->interval = atof(optarg); break;
	case 't': config->task_progress_threshold = atof(optarg); break;
	case 'v':
		dbg_flags = 0;
		flag_token = strtok(optarg, ",");
		if (!strcmp(flag_token, "ALL")) {
			for (enum dbg_flag_t d = DSTARTUP; d < DEND_FLAGS; d++) dbg_flags |= (1 << d);
			break;
		}        
		do {
			for (enum dbg_flag_t d = DSTARTUP; d < DEND_FLAGS; d++) {
				if (!strcmp(dbg_names[d], flag_token)) dbg_flags |= (1 << d);
			}
		} while ((flag_token = strtok(NULL, ",")) != NULL);
		break;
	case 's': config->init_time = atof(optarg); break;
	case 'u': config->enable_numa_migrate = 1; break;
	case 'b': 
		if (strncmp(optarg, "speed", 7) == 0) config->balancing = BALANCE_SPEED;
		else if (strncmp(optarg, "load", 7) == 0) config->balancing = BALANCE_LOAD;
		else if (strncmp(optarg, "static", 7) == 0) config->balancing = BALANCE_STATIC;
		else return ARGP_ERR_UNKNOWN;
		break;
	case 'e': config->sanity_check = 1; break;
	case 'p': config->use_cpu_idle = 1; break;
	case 'r': config->add_randomness = 1; break;
	case 'I': config->expected_idle_task = atoi(optarg); break;
	case 'f': strcpy(config->pid_fname, optarg); break;
	#ifdef USE_NETLINK
	case 'l': config->use_netlink = 1; break;
	#endif
	case ARGP_KEY_END:
		if (IS_EMPTY(config->pid_fname) && IS_EMPTY(config->exec_name)) argp_usage(state);
		break;
	case ARGP_KEY_ARG:
		if (!IS_EMPTY(config->pid_fname)) break;
		strcpy(config->exec_name, &state->argv[state->next - 1][0]);
		config->exec_args = &state->argv[state->next - 1];
		state->next = state->argc;
		break;
	default:
		return ARGP_ERR_UNKNOWN;
	}
	return 0;
}

void get_config(int argc, char** argv, config_t* config) {
	static char dbg_flag_str[200] = "verbose flags - ";
	for (enum dbg_flag_t d = DSTARTUP; d < DEND_FLAGS; d++) {
		strcat(dbg_flag_str, dbg_names[d]);
		if (d != DEND_FLAGS) strcat(dbg_flag_str, ",");
	}
	strcat(dbg_flag_str, "ALL");
	static struct argp_option options[] = {
		{"cores", 'c', "char*", 0, "string representing cores, e.g. 0,1,3 or 4-8,11"},
		{"tasks", 'n', "int", 0, "total number of tasks expected (including any idle task)"},
		{"interval", 'i', "double", 0, "balancing interval in seconds"},
		{"threshold", 't', "double", 0, "progress below average required to classify a task "
		 "as behind (-1 means automatic)"},
		{"verbose", 'v', "char*", 0, dbg_flag_str},
		{"inittime", 's', "double", 0, "initial period spent polling for tasks before balancing"},
		{"numa", 'u', 0, 0, "enable cross NUMA migrations"},
		{"balancing", 'b', "char*", 0, "balancing, 'speed', 'load' or 'static' (defaults to 'speed')"},
		{"sanity", 'e', 0, 0, "sanity check (performs various expensive checks for errors)"},
		{"cpuidle", 'p', 0, 0, "use CPU idle time - only applies to apps that sleep. Needs kernel module"},
		{"addrnd", 'r', 0, 0, "add randomness to balance interval (up to 20%)"},
		{"idletask", 'I', "int", 0, "specify the index of the idle task for apps that have a task that "
		 "launches and then sleeps. This is usually 1 for OpenMP."},
		{"pidfname", 'f', "char*", 0, "Read pids from file name"},
		#ifdef USE_NETLINK
		{"usenl", 'l', 0, 0, "use efficient netlink interface (may require root)"},
		#endif		
		{0}
	};
	// set defaults for config
	strcpy(config->exec_name, "");
	config->exec_args = NULL;
	strcpy(config->exec_args_str, "");
	config->num_cpus = 0;
	config->interval = 0.1;
	// this default means automatic setting
	config->task_progress_threshold = -1.0;
	config->num_tasks = 1;
	config->init_time = 1.0;
	config->enable_numa_migrate = 0;
	config->balancing = BALANCE_SPEED;
	config->sanity_check = 0;
	config->use_cpu_idle = 0;
	config->add_randomness = 0;
	config->expected_idle_task = NOT_SET;
	strcpy(config->pid_fname, "");
	config->use_netlink = 0;
	// get the command line args
	static struct argp argp = {options, parse_opt, "[cmd args_to_cmd]", "Options:"}; 
	int arg_index = 1;
	argp_parse(&argp, argc, argv, ARGP_IN_ORDER, &arg_index, config);
	// convert the interval into a nanosec structure
	config->interval_ts.tv_sec = (time_t)config->interval;
	config->interval_ts.tv_nsec = (config->interval - (time_t)config->interval) * 1e+9;
	// print the results
	printf(_BLUE_ "Juggle v%s, pid %u\n" _END_, JUGGLE_VERSION, getpid());
	if (config->exec_args) {
		for (int i = 0; config->exec_args[i]; i++) {
			strcat(config->exec_args_str, config->exec_args[i]);
			strcat(config->exec_args_str, " ");
		}
	}
	// this is the command to execute, followed by args for that command
	switch (config->balancing) {
	case BALANCE_LOAD: printf(_BLUE_ "LOAD balancing for: " _END_); break;
	case BALANCE_SPEED: printf(_BLUE_ "Speedbalancing for: " _END_); break;
	case BALANCE_STATIC: printf(_BLUE_ "STATIC for: " _END_); break;
	}
	if (!IS_EMPTY(config->pid_fname)) {
		char hostname[255] = "";
		CHECK_CALL(gethostname(hostname, 254));
		strcat(config->pid_fname, ".");
		strcat(config->pid_fname, hostname);
		printf(_BLUE_ " tids in '%s'\n" _END_, config->pid_fname);
	} else printf(_BLUE_ " '%s'\n" _END_, config->exec_args_str);
	printf(_BLUE_ "Number of tasks expected: %d\n" _END_, config->num_tasks);
	printf(_BLUE_ "Number of cpus available: %d\n" _END_, config->num_cpus);
	printf(_BLUE_ "Cpus: " _END_);
	for (int i = 0; i < config->num_cpus; i++) {
		printf(_BLUE_ "%d" _END_, config->cids[i]);
		if (i < config->num_cpus - 1) printf(_BLUE_ "," _END_);
		else printf(_BLUE_ "\n" _END_);
	}
	if (config->balancing == BALANCE_SPEED) {
		printf(_BLUE_ "Interval: %.3f\n" _END_, config->interval);
		printf(_BLUE_ "Task progress threshold: %.3f" _END_, config->task_progress_threshold);
		if (config->task_progress_threshold == -1) {
			// now we automatically adjust for the expected lag
			// FIXME: this should take NUMA into account
			double R = (float)config->num_tasks / (float)config->num_cpus;
			config->task_progress_threshold = 1.0 / ceil(R) / floor(R);
			//printf(_BLUE_ " <%f> " _END_, config->task_progress_threshold);
			config->task_progress_threshold *= ((double)config->interval / 8.0);
			if (config->task_progress_threshold < 0.001) config->task_progress_threshold = 0.001;
			printf(_BLUE_ " (adjusted to %.3f)\n" _END_, config->task_progress_threshold);
		} else {
			printf(_BLUE_ "\n" _END_);
		}
		printf(_BLUE_ "Initialization time: %.3f\n" _END_, config->init_time);
		if (config->enable_numa_migrate) printf(_BLUE_ "Cross NUMA migrations enabled\n" _END_);
		else printf(_BLUE_ "Cross NUMA migrations disabled\n" _END_);
		if (config->sanity_check) printf(_BLUE_ "Sanity checks enabled (*slow*)\n" _END_);
		if (config->use_cpu_idle) printf(_BLUE_ "Using CPU idle time\n" _END_);
		if (config->add_randomness) printf(_BLUE_ "Adding randomness to balance interval\n" _END_);
		if (config->expected_idle_task != NOT_SET) {
			printf(_BLUE_ "Expecting idle task %d\n" _END_, config->expected_idle_task);
		}
		printf(_BLUE_ "Verbose flags: " _END_);
		for (enum dbg_flag_t d = DSTARTUP; d < DEND_FLAGS; d++) {
			if (dbg_flags & (1 << d)) printf(_BLUE_ "%s " _END_, dbg_names[d]);
		}
	}
	#ifndef USE_NETLINK
	config->use_netlink = 0;
	#endif
	if (config->use_netlink) {
		printf(_BLUE_ "\nUsing netlink taskstats %s" _END_, "interface");
	}
	printf("\n");
}

